{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Small Language Model Conversion and Inference with MLC LLM\n",
        "This notebook is a companion of chapter 10 of the \"Domain Specific LLMs in Action\" book, author Guglielmo Iozzia, [Manning Publications](https://www.manning.com/), 2024.  \n",
        "The code in this notebook shows how to use [MLC LLM](https://llm.mlc.ai/) to convert and compile a Small Language Model hosted in the Hugging Face's Hub and then run inference with it on a Linux system. The model under consideration is of [RedPajama-INCITE-Instruct-3B-v1](https://huggingface.co/togethercomputer/RedPajama-INCITE-Instruct-3B-v1), but the code in this notebook applies to any other Open Source LLM hosted in the HF's Hub. Hardware acceleration is required.   \n",
        "More details about the code can be found in the related book's chapter."
      ],
      "metadata": {
        "id": "C0-RkVmQtz-6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the proper MLC LLM wheel for Linux and the CUDA drivers in this system."
      ],
      "metadata": {
        "id": "BgYuYYc-vwNQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-cu122 mlc-ai-nightly-cu122"
      ],
      "metadata": {
        "id": "JkN8iOWWCJHA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Verify the MLC installation completed successfully."
      ],
      "metadata": {
        "id": "BFTuUzEjwDe1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mlc_llm --help"
      ],
      "metadata": {
        "id": "DsoTroR6CX9J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Model Conversion"
      ],
      "metadata": {
        "id": "C-jkfmpuz3ad"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "To run a model with MLC LLM, we need to convert the model weights into MLC format. Some preliminary actions to be done: create the destination directory for the original model's weights and accessory files, install the Git extension for versioning large files and clone the HF's repo for the target model."
      ],
      "metadata": {
        "id": "3BqlH2QdCzLk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir -p dist/models && cd dist/models\n",
        "!git lfs install\n",
        "!git clone https://huggingface.co/togethercomputer/RedPajama-INCITE-Instruct-3B-v1"
      ],
      "metadata": {
        "id": "kpxINUTEDQBT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Convert the model weights to the MLC LLM format. The converted weights are saved into the same directory as for the original model."
      ],
      "metadata": {
        "id": "E2t1nliGxCPb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mlc_llm convert_weight ./RedPajama-INCITE-Instruct-3B-v1/ \\\n",
        "    --quantization q4f16_1 \\\n",
        "    -o dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC \\\n",
        "    --device cuda:0"
      ],
      "metadata": {
        "id": "RWDDM9TzjZW4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate the chat configuration for the converted model. The generated configuration is saved in the same directory as for the converted weights."
      ],
      "metadata": {
        "id": "rCD918ZpySVu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mlc_llm gen_config ./RedPajama-INCITE-Instruct-3B-v1/ \\\n",
        "    --quantization q4f16_1 --conv-template redpajama_chat \\\n",
        "    -o dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC/"
      ],
      "metadata": {
        "id": "Z50i2An0DXVm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Verify that all the required files (the chat configuration file, the model's weights info, shards and tokenizer files) are within the destination directory."
      ],
      "metadata": {
        "id": "OX51MRVgyjV_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!ls dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC"
      ],
      "metadata": {
        "id": "A9glCK0yDuOr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We need now to compile the converted model before we can run inference. Create a destination directory for the compiled model."
      ],
      "metadata": {
        "id": "Wq-JbWtfzVht"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir ./dist/libs"
      ],
      "metadata": {
        "id": "rz6ido7kKHLP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Compile the model library using the specification in chat configuration file preliminary created (*mlc-chat-config.json*)."
      ],
      "metadata": {
        "id": "tlUEzCMdzkQr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mlc_llm compile ./dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC/mlc-chat-config.json \\\n",
        "    --device cuda -o dist/libs/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-cuda.so"
      ],
      "metadata": {
        "id": "d8u50YXoDvKv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Verify that the model compilation completed successfully. For Linux systems and CUDA drivers the compilation directory should contain a single compiled library file. Please refer to the official MLC LLM documentation for other operating systems and hardware accelerators."
      ],
      "metadata": {
        "id": "oCdrRd350C8N"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!ls dist/libs"
      ],
      "metadata": {
        "id": "RWC5TlhGERV3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Chat with the Converted model using the MLC LLM Python API"
      ],
      "metadata": {
        "id": "IOgCLquyz7F4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Create an instace of the MLCEngine for the converted model. This class supports only synchronous chat completions."
      ],
      "metadata": {
        "id": "fjyy3R3d0QBc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from mlc_llm import MLCEngine\n",
        "\n",
        "engine = MLCEngine(model=\"./dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC\",\n",
        "                   model_lib=\"./dist/libs/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-cuda.so\")"
      ],
      "metadata": {
        "id": "-yCed4zaESIF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Start some chat examples."
      ],
      "metadata": {
        "id": "M1D69iOv2GEa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for response in engine.chat.completions.create(\n",
        "    messages=[{\"role\": \"user\", \"content\": \"What's the meaning of life?\"}],\n",
        "    model=\"./dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC\",\n",
        "    stream=False,\n",
        "):\n",
        "    print(response)\n",
        "print(\"\\n\")"
      ],
      "metadata": {
        "id": "Xb45pBf0OALS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for response in engine.chat.completions.create(\n",
        "    messages=[{\"role\": \"user\", \"content\": \"What's the meaning of life?\"}],\n",
        "    model=\"./dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC\",\n",
        "    stream=True,\n",
        "):\n",
        "    for choice in response.choices:\n",
        "        print(choice.delta.content, end=\"\", flush=True)\n",
        "print(\"\\n\")"
      ],
      "metadata": {
        "id": "D8IyWurSGmd2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Shutdown the MLC engine."
      ],
      "metadata": {
        "id": "zjswzwmX2BEX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "engine.terminate()"
      ],
      "metadata": {
        "id": "hMG7TiopNh6d"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}